import os
import numpy as np
from scipy.interpolate import griddata
import re

def process_cas_files(input_file_path, output_file_path):
    """
    处理指定文件夹下的.cas文件，筛选符合要求的段落保存到新文件
    :param folder_path: 存放.cas文件的文件夹路径
    :param output_file_path: 输出文件的路径
    """
    with open(output_file_path, 'w') as output_file:
        with open(input_file_path, 'r') as input_file:
            current_paragraph = []
            in_section = False
            for line in input_file:
                line = line.strip()
                if (line.startswith('(10 ') or line.startswith('(13 ')) and not line.endswith('))'):
                    in_section = True
                    current_paragraph = []
                    current_paragraph.append(line)
                elif line.endswith('))') and in_section:
                    in_section = False
                    current_paragraph.append(line)
                    for section_line in current_paragraph:
                            output_file.write(section_line + '\n')
                    output_file.write('\n')
                elif in_section:
                    current_paragraph.append(line)
                else:
                    continue
            

def process_dat_files(input_file_path, output_file_path):
    """
    处理指定文件夹下的.dat文件，筛选符合要求的段落保存到新文件
    :param folder_path: 存放.dat文件的文件夹路径
    :param output_file_path: 输出文件的路径
    """
    with open(output_file_path, 'w') as output_file:
        with open(input_file_path, 'r') as input_file:
            current_paragraph = []
            for line in input_file:
                line = line.strip()
                if line.startswith('(0'):
                    # 遇到分隔行，检查当前段落第一行是否符合要求
                    if current_paragraph:
                        first_line = current_paragraph[0].lower()
                        has_cell_zone_20 = "cell zone 20" in first_line
                        has_keyword = any(keyword in first_line for keyword in ["sv_p,", "sv_u,", "sv_v,", "sv_w,", "x-coordinate", "y-coordinate", "z-coordinate"])
                        if has_cell_zone_20 and has_keyword:
                            output_file.write('\n'.join(current_paragraph) + '\n')
                    current_paragraph = []
                    current_paragraph.append(line)
                else:
                    current_paragraph.append(line)
            # 处理文件末尾的段落
            if current_paragraph:
                first_line = current_paragraph[0].lower()
                has_cell_zone_20 = "cell zone 20" in first_line
                has_keyword = any(keyword in first_line for keyword in ["sv_p,", "sv_u,", "sv_v,", "sv_w,", "x-coordinate", "y-coordinate", "z-coordinate"])
                if has_cell_zone_20 and has_keyword:
                    output_file.write('\n'.join(current_paragraph) + '\n')

def get_coordinates(dat_file_path, dimension):
    with open(dat_file_path, 'r') as file:
        coordinates = []
        in_section = False
        for line in file:
            line = line.strip()
            if line.startswith('(10'):
                in_section = True
                continue
            if in_section and line.endswith('))'):
                point = [float(num) for num in line[:-2].split()]
                coordinates.append(point)
                in_section = False
                break
            if in_section and line != '(':
                 point = [float(num) for num in line.split()]
                 if len(point) == dimension:
                    coordinates.append(point)
    return coordinates

def get_boundary_node(dat_file_path, coordinate, dimension):
    mesh = {}
    mesh_coord = {}
    in_section = False

    mesh['interior'] = []
    mesh_coord['interior'] = []

    with open(dat_file_path, 'r') as file:
        for line in file:
            line = line.strip()
            if line.startswith('(13'):
                parts = line.split()
                
                if int(parts[4], 16) == 2:
                    continue
                if int(parts[4], 16) == 3:
                    in_section = 'wall'
                if int(parts[4], 16) == 5:
                    in_section = 'pressure_outlet'
                if int(parts[4], 16) == 10:
                    in_section = 'velocity_inlet'
                for index in range(100):
                    if ((in_section + '_' + str(index)) not in mesh):
                        in_section = in_section + '_' + str(index)
                        break
                continue

            if in_section:
                if line.endswith('))'):
                    parts = line[:-2].split()
                elif line != '(':
                    parts = line.split()
                else:
                    continue
                
                if in_section not in mesh:
                    mesh[in_section] = []
                    mesh_coord[in_section] = []
                if dimension == 3:
                    for index_str in parts[1 : int(parts[0]) + 1]:
                        mesh[in_section].append(int(index_str, 16))
                else:
                    for index_str in parts[0 : 2]:
                        mesh[in_section].append(int(index_str, 16))
                if line.endswith('))'):
                    in_section = 0
        

        for i, coord in enumerate(coordinate):
            is_interior = True
            for boundary_name, boundary_index in mesh.items():
                if (i + 1) in boundary_index:
                    mesh_coord[boundary_name].append(coord)
                    is_interior = False
                    break
            if is_interior:
                mesh_coord['interior'].append(coord)
                
    return mesh_coord

def extract_data_to_matrix(dat_file_path):
    """
    从给定的.dat文件中提取不同类型的数据并整合到一个numpy矩阵中
    :param dat_file_path: 之前生成的.dat文件路径
    :return: 整合好数据的numpy矩阵
    """
    data_dict = {
        "x-coordinate": [],
        "y-coordinate": [],
        "z-coordinate": [],
        "SV_U": [],
        "SV_V": [],
        "SV_W": [],
        "SV_P": []
    }
    with open(dat_file_path, 'r') as input_file:
        current_paragraph = []
        current_keyword = None
        for line in input_file:
            line = line.strip()
            if line.startswith('(0'):
                if current_keyword:
                    if len(current_paragraph) >= 4:
                        data_lines = current_paragraph[3:-2]
                        if current_keyword == "SV_U":
                            data_dict["SV_U"].extend([float(x) for x in data_lines])
                        elif current_keyword == "SV_V":
                            data_dict["SV_V"].extend([float(x) for x in data_lines])
                        elif current_keyword == "SV_W":
                            data_dict["SV_W"].extend([float(x) for x in data_lines])
                        elif current_keyword == "SV_P":
                            data_dict["SV_P"].extend([float(x) for x in data_lines])
                        elif current_keyword == "x-coordinate":
                            data_dict["x-coordinate"].extend([float(x) for x in data_lines])
                        elif current_keyword == "y-coordinate":
                            data_dict["y-coordinate"].extend([float(x) for x in data_lines])
                        elif current_keyword == "z-coordinate":
                            data_dict["z-coordinate"].extend([float(x) for x in data_lines])

                if current_paragraph:
                    first_line = current_paragraph[0].lower()
                for keyword in data_dict.keys():
                    if keyword in line:
                        current_keyword = keyword
                        break
                
                current_paragraph = []
                current_paragraph.append(line)
            else:
                current_paragraph.append(line)
        
        if current_paragraph:
            data_lines = current_paragraph[3:-2]
            if current_keyword == "SV_U":
                data_dict["SV_U"].extend([float(x) for x in data_lines])
            elif current_keyword == "SV_V":
                data_dict["SV_V"].extend([float(x) for x in data_lines])
            elif current_keyword == "SV_W":
                data_dict["SV_W"].extend([float(x) for x in data_lines])
            elif current_keyword == "SV_P":
                data_dict["SV_P"].extend([float(x) for x in data_lines])
            elif current_keyword == "x-coordinate":
                data_dict["x-coordinate"].extend([float(x) for x in data_lines])
            elif current_keyword == "y-coordinate":
                data_dict["y-coordinate"].extend([float(x) for x in data_lines])
            elif current_keyword == "z-coordinate":
                data_dict["z-coordinate"].extend([float(x) for x in data_lines])

    # 将各类型数据整理成矩阵（确保各数据长度相同，这里简单以第一个数据类型长度为准来截断）
    min_length = min(len(v) for v in data_dict.values() if v)
    matrix_data = []
    for keyword in data_dict.keys():
        matrix_data.append(data_dict[keyword][:min_length])
    matrix_data = np.array(matrix_data).T
    x = matrix_data[:, 0]
    y = matrix_data[:, 1]
    z = matrix_data[:, 2]
    # x_min, x_max = np.min(x), np.max(x)
    # y_min, y_max = np.min(y), np.max(y)
    # z_min, z_max = np.min(z), np.max(z)
    # half_x = x_min + (x_max - x_min) / 14 #(x_min + x_max) / 2
    # half_y = y_min + (y_max - y_min) / 4 #(y_min + y_max) / 2
    # half_z = z_min + (z_max - z_min) / 2 #(z_min + z_max) / 2
    x_min, x_max = -1.5, 0.5
    y_min, y_max = -1, 1
    z_min, z_max = 0, 2
    mask = (x <= x_max) & (x >= x_min) & (y <= y_max) & (y >= y_min) & (z <= z_max) & (z >= z_min)
    matrix_data = matrix_data[mask]
    x = matrix_data[:, 0]
    y = matrix_data[:, 1]
    z = matrix_data[:, 2]
    sv_u_min, sv_u_max = -18, 63
    sv_v_min, sv_v_max = -30, 43
    sv_w_min, sv_w_max = -27, 36
    sv_p_min, sv_p_max = -2600, 1100 
    sv_u = (matrix_data[:, 3] - sv_u_min) / (sv_u_max - sv_u_min)
    sv_v = (matrix_data[:, 4] - sv_v_min) / (sv_v_max - sv_v_min)
    sv_w = (matrix_data[:, 5] - sv_w_min) / (sv_w_max - sv_w_min)
    sv_p = (matrix_data[:, 6] - sv_p_min) / (sv_p_max - sv_p_min)

    # 找到 y >= 0 的数据的索引
    positive_y_indices = y >= 0
    # 镜像 x, z 坐标和 SV_P, SV_U, SV_W 数据
    matrix_data = []
    matrix_data.append(np.concatenate((x[positive_y_indices], x[positive_y_indices])))
    matrix_data.append(np.concatenate((y[positive_y_indices], -y[positive_y_indices])))
    matrix_data.append(np.concatenate((z[positive_y_indices], z[positive_y_indices])))
    matrix_data.append(np.concatenate((sv_u[positive_y_indices], sv_u[positive_y_indices])))
    matrix_data.append(np.concatenate((sv_v[positive_y_indices], -sv_v[positive_y_indices])))
    matrix_data.append(np.concatenate((sv_w[positive_y_indices], sv_w[positive_y_indices])))   
    matrix_data.append(np.concatenate((sv_p[positive_y_indices], sv_p[positive_y_indices])))
    matrix_data = np.array(matrix_data) 

    return matrix_data.T

def scale_and_interpolate(matrix_data, scale_length):
    """
    根据坐标确定盒子、缩放、构建新坐标系并进行反距离插值
    :param matrix_data: 从.dat文件提取并整合好的numpy矩阵
    :param scale_length: 向内缩放的长度
    :return: 插值后的速度矩阵（Vel_matrix）和压力矩阵（P_matrix）
    """
    x = matrix_data[:, 0]
    y = matrix_data[:, 1]
    z = matrix_data[:, 2]
    x_min, x_max = np.min(x), np.max(x)
    y_min, y_max = np.min(y), np.max(y)
    z_min, z_max = np.min(z), np.max(z)
    # 确定缩放后的盒子范围
    x_box_min = x_min + scale_length
    x_box_max = x_max - scale_length
    y_box_min = y_min + scale_length
    y_box_max = y_max - scale_length
    z_box_min = z_min + scale_length
    z_box_max = z_max - scale_length
    # 构建新的正交坐标系（这里简单示例，可根据实际需求调整分辨率等）
    x_interval = (x_box_max - x_box_min) / 40  # 这里先初始化为一个参考值，后续会调整
    y_interval = (y_box_max - y_box_min) / 40
    z_interval = (z_box_max - z_box_min) / 40
    interval = min(x_interval, y_interval, z_interval)
    new_x = np.arange(x_box_min, x_box_max + interval / 2, interval)
    new_y = np.arange(y_box_min, y_box_max + interval / 2, interval)
    new_z = np.arange(z_box_min, z_box_max + interval / 2, interval)
    # 线性插值
    points = matrix_data[:,:3]
    values = matrix_data[:,3:]
    grid_points = np.meshgrid(new_x, new_y, new_z)
    grid_points = np.array(grid_points).reshape(3,-1).T
    interpolated_data = griddata(points, values, grid_points, method='linear').T
    # 构建4阶张量
    VelX_matrix = interpolated_data[0].reshape(new_x.size, new_y.size, new_z.size)
    VelY_matrix = interpolated_data[1].reshape(new_x.size, new_y.size, new_z.size)
    VelZ_matrix = interpolated_data[2].reshape(new_x.size, new_y.size, new_z.size)
    Vel_matrix = np.stack([VelX_matrix, VelY_matrix, VelZ_matrix], axis=0)
    P_matrix = interpolated_data[3].reshape(1, new_x.size, new_y.size, new_z.size)
    
    return Vel_matrix, P_matrix

# 使用示例，你需要替换为实际的文件夹路径和输出文件路径
folder_path = "D:\\yanlin.ren\\proceeding_job\\20250313pipetee2D"
scale_length = 0.1

def main():
    dimension = 2
    files = [file for file in os.listdir(folder_path) if file.endswith('.cas')]
    for file in files:
        input_file_path = os.path.join(folder_path, file)
        output_file_path = os.path.join(folder_path, "temp.txt")
        process_cas_files(input_file_path, output_file_path)
        coordinates = get_coordinates(output_file_path, dimension)
        mesh_coord = get_boundary_node(output_file_path, coordinates, dimension)

        np.savez(os.path.join(folder_path, f'mesh.npz'), **mesh_coord)


if __name__ == "__main__":
    main()